import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed
from dify_api import dify_api


class CSVProcessor:
    def __init__(self, input_csv_path, output_csv_path, max_workers=10):
        self.input_csv_path = input_csv_path
        self.output_csv_path = output_csv_path
        self.api_query = dify_api.DifyAPIQuery()
        self.max_workers = max_workers

    def process_row(self, index, row):
        try:
            user_question = row['content']
            user_id = row['user_id']
            print(f"Processing row {index}: {user_question}")

            # 调用 API 获取回复内容
            response_content = self.api_query.send_query(user_question, user_id)
            print(f"Reply for row {index}: {response_content['answer']}")

            return index, response_content['answer']
        except Exception as e:
            print(f"Error processing row {index}: {e}")
            return index, None

    def process_csv_file(self):
        # 读取 CSV 文件
        df = pd.read_csv(self.input_csv_path)

        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_index = {
                executor.submit(self.process_row, index, row): index
                for index, row in df.iterrows()
            }

            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    index, reply = future.result()
                    # 将回复内容存储到新列 intelligent_reply 中
                    df.at[index, 'intelligent_reply'] = reply
                except Exception as e:
                    print(f"Error in future for row {index}: {e}")

        # 将结果保存到新的 CSV 文件
        df.to_csv(self.output_csv_path, index=False)

# 示例使用
input_csv_path = 'chat_messages.csv'  # 输入文件路径
output_csv_path = 'batch_test_result.csv'  # 输出文件路径
processor = CSVProcessor(input_csv_path, output_csv_path, max_workers=30)
processor.process_csv_file()